package dbAnnoAndIndexCreate;

use strict;
use File::Basename qw(basename dirname);
use Cwd qq(abs_path);
use Carp;

require Exporter;

our @ISA         = qw(Exporter);
our @EXPORT      = qw(annoByIndex);
our @EXPORT_OK   = qw(makeIndex);
our $VERSION     = 0.1;

=head2 annoByIndex

    Title   : annoByIndex
    Function: 
    Usage   : &annoByIndex($dbFile,$queryfile,$annotatedfile)
    Returns : 
    Args    : file

=cut
sub annoByIndex
{

    confess "annoByIndex(dbFile,queryfile,outAnnotatedFile) needs three args" unless scalar(@_) ==3;
    my $dbFile = shift;
    my $queryfile = shift;
    my $annotatedfile = shift;

    my $indexfilter_threshold = 0.9;

    sub printerr;
    &filterQueryThread($dbFile,$queryfile,$annotatedfile,$indexfilter_threshold);

}

sub filterQueryThread 
{
    my ($dbFile,$queryfile,$outfile,$indexfilter_threshold) = @_;
    
    my ($FIL, $DROPPED, $INVALID,$title);
    #open ($FIL, ">$filtered_file") or die "Error: cannot write to output file $outfile: $!\n"; 
    open ($DROPPED, ">$outfile") or die "Error: cannot write to output file $outfile: $!\n";
    #open ($INVALID, ">$invalid_file") or die "Error: cannot write to output file $invalid_file: $!\n";
            
    open (QUERY, $queryfile) or die "Error: cannot read from query file $queryfile: $!\n";
    
    my (%variant, $filedone, $batchdone);
    my ($linecount, $batchlinecount, $invalid, $invalidcount) = (0, 0);
    my ($chr, $start, $end, $ref, $obs, $info);
    
    
    while (1) {
        $_ = <QUERY>;
      
        if (not defined $_) {
            $filedone++;
        } else {
            
            if ($_ =~ /Zygosity/){

                $title = $_;
                chomp($title);
                print $DROPPED "$title\tClinPred_score\tClinPred_pred\n";
                next;
            }
            #$_ =~ /Zygosity/ and $title = $_ and print DROPPED  and next;
            
            s/[\r\n]+$//;
            
            $linecount++;
            $batchlinecount++;

            #Chr     Pos     Ref     Alt     Zygosity
            #NC_000001       69270   A       G       hom
            #($invalid, $chr, $start, $end, $ref, $obs) = detectInvalidInput ($_);
            ($chr, $start, $ref, $obs) = split /\s+/,$_;
            $chr = &NC2Chr($chr) if $chr =~ /^NC/i;#2020-02-02
            $chr =~ s/^chr//i if $chr =~ /^chr/i;#2020-02-02

            if ( $obs =~ /\// ){
                $obs =~ s/.*\///g;
            }            

          
            if (exists $variant{$chr, $start, $ref,$obs}) {
                $variant{$chr, $start, $ref,$obs} .= "\n$_";
            } else {
                $variant{$chr, $start, $ref,$obs} = "$_";  #store the information for the variant in newline-separated strings, with the first field being ref allele
            }
        }
        
        if ($filedone) {
            %variant and printerr "NOTICE: Processing next batch with ${\(scalar keys %variant)} unique variants in $batchlinecount input lines\n";
            %variant and filterNextBatchThread ($dbFile,\%variant, $FIL, $DROPPED, $INVALID,$indexfilter_threshold);       #filter the next batch of variants (print output to FIL, DROPPED), then clear %variant hash for next batch
            %variant = ();
            $batchlinecount = 0;                #reset the line count for this batch
            $batchdone = 0;
        }
        if ($filedone) {
            #close ($FIL);
            close ($DROPPED);
            #close ($INVALID);
            last;
        }
    }
}

sub filterNextBatchThread 
{
    my ($dbfile,$variant, $FIL, $DROPPED, $INVALID,$indexfilter_threshold) = @_;
#    my $dbfile = "/home/wangshenchang/database/humandb/ClinPred_v2.txt";
    my $dbtype1 = '';
    my $dbtype = '';
    
    -f $dbfile or die "Error: the database file $dbfile is not present\n";
    
    my (@record, $chr, $start, $end, $ref, $obs, $score, @otherinfo, $qual, $fil, $info);       #@otherinfo: other information about the variant in addition to score (by default, only score for the variant will be printed out)
    my ($rsid, $strand, $ucscallele, $twoallele, $class, $af, $attribute);
    #my $count_invalid_dbline;

    my ($BIN, $DBSIZE, $NUMBIN) = (0, 0);
    my %index = ();         #holds all bin index, used by DBM only and not by IDX (20140725)
    my $bb = {};            #a subset of %index, which corresponds to the input variants
    my $flag_idx_search = 0;    #indicate if index-based search algorithm is used (faster speed for a small number of input variants)
    

    if (-f "$dbfile.idx") {
        
        if (open(IDX, "$dbfile.idx")) {
            my $line = <IDX>;
            if (not $line =~ m/BIN\t(\d+)\t(\d+)/) {
                printerr "WARNING: Malformed database index file $dbfile.idx.\n";
            } elsif ($2 != -s $dbfile) {            #file version is different, do not use index file in this case
                printerr "WARNING: Your index file $dbfile.idx is out of date and will not be used. ANNOVAR can still generate correct results without index file.\n";
            } else {
                ($BIN, $DBSIZE) = ($1, $2);

                foreach my $k ( keys %$variant ) {
                    my ($chrom, $pos) = split ($;, $k);
                    
                    my $bin = $pos - ($pos % $BIN);
                    
                    $bb->{"$chrom\t$bin"} = 0;  #flag this bin to be searched later
                }               
                
                my ($count_total, $count_search) = qw/0 0/;
                while ( $line = <IDX> ) {
                    $line =~ s/[\r\n]+$//;
                    my ( $chrom, $bin, $offset0, $offset1 ) = split (/\t/, $line);
                    $chrom =~ s/^chr//;     #delete the chr in snp135, etc 
                    defined $offset1 or next;   #invalid input line in the index file
                    
                    if (defined $bb->{"$chrom\t$bin"}) {
                        $bb->{"$chrom\t$bin"} or $count_search++;
                        $bb->{"$chrom\t$bin"} = "$offset0,$offset1";
                    }
                    $count_total++;
                }
                if ($count_total and $count_search/$count_total < $indexfilter_threshold) {
                    $flag_idx_search++;
                    printerr "NOTICE: Database index loaded. Total number of bins is $count_total and the number of bins to be scanned is $count_search\n";
                }
            }
            close (IDX);
        } else {
            printerr "WARNING: Failed to read from filter database index $dbfile.idx: $!. Switching to text scanning-based filter operation\n";
        }
    }

    if (not $flag_idx_search) {
        $bb = {1, join (',', 0, -s "$dbfile")};
        %index = (1, join (',', 0, -s "$dbfile"));
    }

    open (DB, $dbfile) or die "Error: cannot read from input database file $dbfile: $!\n";
    printerr "NOTICE: Scanning filter database $dbfile...";
    
    foreach my $b (sort keys %$bb) {
        
        $bb->{$b} or next;                  #DB does not have this bin to search against (value is zero)
        
        my ($chunk_min, $chunk_max) = split (/,/, $bb->{$b});
        
        defined $chunk_max or die "Error: b=$b index=$index{$b} bb=$bb->{$b}\n";

        seek(DB, $chunk_min, 0);                #place file pointer to the chunk_min
        my $chunk_here = $chunk_min;

        while (<DB>) {
            my $line_length = length($_);           #calculate line length of the DB

            $chunk_here += $line_length;
    
            my (@obs2, @score2, @start2, @ref2);                #for 1000G2010 data set in VCF format, some tri-allelic SNPs are present; in the future, some quad-allelic SNPs may be also present in VCF files
            s/[\r\n]+$//;
            m/\S/ or next;                  #skip empty lines in the database file (sometimes this occurs)
            m/^#/ and next;                 #skip the comment line
 
             {
                #$dbtype eq 'generic' or print STDERR "NOTICE: the --dbtype $dbtype is assumed to be in generic ANNOVAR database format\n";
                #($chr, $start, $end, $ref, $obs, $score, @otherinfo) = split (/\t/, $_);
                ($chr, $start, $end, $ref, $obs, $score, @otherinfo) = split (/\t/, $_);
                defined $obs or die "Error: invalid database entry (generic ANNOVAR database should contain at least 4 tab-delimited fields per line): <$_>\n";
                ($ref, $obs) = (uc $ref, uc $obs);      #make sure to use upper case, as query is always in upper case
                defined $obs or die "Error: the generic database file must contains at least five tab-delimited fields per line (but observed line: $_)\n";
                #defined $score or $score = "NA";
                defined $score or $score = 0;
                $chr =~ s/^chr//i;      #changed 20120712; when genericdb contains "chr", all variants will be filtered.

                defined $obs or die "Error: invalid record found in DB file $dbfile (at least 5 fields expected for 'generic' dbtype): <$_>\n";

            }

            if ($variant->{$chr, $start, $ref,$obs}) {
                my (@info) = split (/\n/, $variant->{$chr, $start, $ref,$obs});    #most likely, only one piece of information
                for my $i (0 .. @info-1) {

                    #print $DROPPED join ("\t", $dbtype, $score), "\t", $info[$i], "\n";
                    $score >0.5 and  print $DROPPED $info[$i], "\t", $score, "\t",  "Deleterious", "\n";
                    $score <=0.5 and  print $DROPPED $info[$i], "\t", $score, "\t",  "Tolerated", "\n";

                }
                delete $variant->{$chr, $start, $ref,$obs};
            }
    
            if ( $chunk_here > $chunk_max ) {
                last;
            }
        }
    }

     for my $key (keys %$variant) {
         my ($chr, $start, $ref,$obs) = split ($;, $key);     #hash key separator
         my (@info) = split (/\n/, $variant->{$key});

         for my $i (0 .. @info-1) {
             #print $FIL $info[$i], "\n";
             print $DROPPED $info[$i], "\t", "-", "\t", "-", "\n";
         }
     }
    printerr "Done\n";
    #$count_invalid_dbline and printerr "WARNING: $count_invalid_dbline lines in dbfile $dbfile were ignored due to invalid formats\n";
}

sub detectInvalidInput 
{
    my ($line) = @_;
    my $invalid = 0;
    my @avcolumn = (0..4);
    my @nextline = split (/\s+/, $line);
    my ($chr, $start, $end, $ref, $obs) = @nextline[@avcolumn];
    if ( not (defined $chr and defined $start and defined $end and defined $ref and defined $obs)) {
        $invalid++;
    } else {
        ($ref, $obs) = (uc $ref, uc $obs);
        #$zerostart and $start++;
        $chr =~ s/^chr//;
        if ($chr =~ m/[^\w\.]/ or $start =~ m/[^\d]/ or $end =~ m/[^\d]/) {     #chr name could contain . (example: GL000212.1, or Zv9_NA###
            $invalid++;
        } elsif ($ref eq '-' and $obs eq '-'        #both are empty allele
            or $ref =~ m/[^ACTG0\-]/        #non-standard nucleotide code
            or $obs =~ m/[^ACGT0\-]/        #non-standard nucleotide code
            or $start =~ m/[^\d]/           #start is not a number
            or $end =~ m/[^\d]/             #end is not a number
            or $start > $end            #start is more than end
            #or $ref ne '0' and $end-$start+1 != length ($ref)   #length mismatch with ref
            #or $ref eq '-' and $start != $end   #length mismatch for insertion
            ) {
            $invalid++;
        }
    }
    return ($invalid, $chr, $start, $end, $ref, $obs);
}

sub printerr 
{
    print STDERR @_;
 #   print LOG @_;
}

sub NC2Chr 
{
    # body...
    my ( $NC_Num ) = @_;
    if( $NC_Num eq 'NC_000023' ){
        return 'X';
    }elsif( $NC_Num eq 'NC_000024' ){
        return 'Y';
    }elsif( $NC_Num =~ m/NC_00000(\d)/ ){
        return "$1";
    }elsif( $NC_Num =~ m/NC_0000(\d+)/ ){
        return "$1";
    }else{
        print STDERR "ERROR:\tChromosomal format is incorrect...\n";
    }
}


=head2 makeIndex

    Title   : makeIndex
    Function: 
    Usage   : &makeIndex($dataFile)
    Returns : 
    Args    : file

=cut
sub makeIndex
{
    confess "makeIndex(dbFile,binSize) needs two args" unless scalar(@_) ==2;
    my $db = shift;
    my $bin_size = shift;

    my $db_absPath = abs_path($db);
    my $db_dir     = dirname($db_absPath);
    my $db_index  = basename($db_absPath) . ".idx";


    if (!-e $db) {
      die "$db not found\n";
    }

    my $file_size = -s $db;

    my %index;
    open(my $in, "<", $db) or die "Couldn't open $db for indexing\n";
    open IDX, ">", $db_dir.$db_index;

    my $previous_file_position = tell $in;
    #print $previous_file_position;

    while (my $ln = <$in>) {
           #Check input file. Some are (chr,start,stop) and others are (id,chr,start,stop).

        next if $ln =~/^#/;
           my ($chr,$start,$stop) = (split "\t", $ln)[0,1,2];
           my $bin_start = int($start/$bin_size) * $bin_size;
           my $current_file_position = tell $in;
           if (!exists $index{$chr}->{$bin_start}) {
                $index{$chr}->{$bin_start} = [$previous_file_position, $current_file_position];
           }
           else{
                $index{$chr}->{$bin_start}->[1] = $current_file_position;
           }
           $previous_file_position = $current_file_position;
    }

    close $in;

    print IDX "#BIN\t$bin_size\t$file_size\n";
    foreach my $chr ((1,10..19,2,20,21,22,3..9,"MT","X","Y")){ 
           #Ordered array to match other Annovar idx files
           my $chr11=$chr;
           foreach my $index_region (sort keys %{$index{$chr11}}){
                  my $start = $index{$chr11}->{$index_region}->[0];
                  my $stop = $index{$chr11}->{$index_region}->[1];
                  print IDX "$chr11\t$index_region\t$start\t$stop\n";
            }
    }
    close IDX;

    return 0;
}


1;


__END__
